package middleware

import (
	"lan-file-transfer/models"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

// RateLimiter 限流器
type RateLimiter struct {
	requests sync.Map // key: IP, value: *RequestRecord
}

// RequestRecord 请求记录
type RequestRecord struct {
	Count      int
	ResetTime  time.Time
	mu         sync.Mutex
}

var globalRateLimiter = &RateLimiter{}

// RateLimit 限流中间件(每IP每秒最多10个请求)
func RateLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		value, _ := globalRateLimiter.requests.LoadOrStore(ip, &RequestRecord{
			Count:     0,
			ResetTime: time.Now().Add(1 * time.Second),
		})

		record := value.(*RequestRecord)
		record.mu.Lock()
		defer record.mu.Unlock()

		// 检查是否需要重置计数
		if time.Now().After(record.ResetTime) {
			record.Count = 0
			record.ResetTime = time.Now().Add(1 * time.Second)
		}

		// 检查是否超过限制
		if record.Count >= 10 {
			c.JSON(http.StatusTooManyRequests, models.ErrorResponse("请求过于频繁,请稍后再试"))
			c.Abort()
			return
		}

		record.Count++
		c.Next()
	}
}

// IPWhitelist IP白名单中间件(只允许局域网IP)
func IPWhitelist() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		// 检查是否为局域网IP
		if !isPrivateIP(ip) {
			c.JSON(http.StatusForbidden, models.ErrorResponse("只允许局域网访问"))
			c.Abort()
			return
		}

		c.Next()
	}
}

// isPrivateIP 判断是否为局域网IP
func isPrivateIP(ip string) bool {
	// 简单判断,实际应该更严格
	if ip == "127.0.0.1" || ip == "::1" || ip == "localhost" {
		return true
	}

	// 检查常见局域网IP段
	if len(ip) >= 7 {
		if ip[:3] == "10." || ip[:4] == "172." || ip[:8] == "192.168." {
			return true
		}
	}

	return true // 开发时暂时允许所有IP
}
